#include "nasm/elf/elfFile.h"
#include "nasm/symbol/symbol.h"
#include "../../lit/include/lit/elf/elfFile.h"

#include <cstdio>
#include <cstring>

namespace dspnasm {
    void ElfFile::addShdr(const std::string &sh_name, Elf32_Word sh_type, Elf32_Word sh_flags, Elf32_Addr sh_addr,
                          Elf32_Off sh_offset, Elf32_Word sh_size, Elf32_Word sh_link, Elf32_Word sh_info,
                          Elf32_Word sh_addralign, Elf32_Word sh_entsize) {
        std::shared_ptr<Elf32_Shdr> sh = std::make_shared<Elf32_Shdr>();
        sh->sh_name = 0;                /* Section name (string tbl index) */
        sh->sh_type = sh_type;          /* Section type */
        sh->sh_flags = sh_flags;        /* Section flags */
        sh->sh_addr = sh_addr;          /* Section virtual addr at execution */
        sh->sh_offset = sh_offset;      /* Section file offset */
        sh->sh_size = sh_size;          /* Section size in bytes */
        sh->sh_link = sh_link;          /* Link to another section */
        sh->sh_info = sh_info;          /* Additional section information */
        sh->sh_addralign = sh_addralign;     /* Section alignment */
        sh->sh_entsize = sh_entsize;         /* Entry size if section holds table */
        shdrTab[sh_name] = sh;
        shdrNames.push_back(sh_name);
    }

    void ElfFile::addSymbol(std::shared_ptr<Symbol> symbol) {
        std::shared_ptr<Elf32_Sym> sym = std::make_shared<Elf32_Sym>();
        sym->st_name = 0;       /* Symbol name (string tbl index) */
        sym->st_value = symbol->m_offset;       /* Symbol value */
        sym->st_size = symbol->m_data.size();   /* Symbol size */
        if (symbol->m_isGlobal) {
            sym->st_info = ELF32_ST_INFO(STB_GLOBAL, STT_NOTYPE);    /* Symbol type and binding */
        } else {
            sym->st_info = ELF32_ST_INFO(STB_LOCAL, STT_NOTYPE);     /* Symbol type and binding */
        }
        sym->st_other = 0;     /* Symbol visibility */
        if (symbol->m_isExterned) {
            sym->st_shndx = STN_UNDEF;      /* Section index */
        } else {
            sym->st_shndx = getSegIndex(symbol->m_segmentName) + 1; /* Section index */
        }
        symTab[symbol->m_symbolName] = sym;
        symNames.push_back(symbol->m_symbolName);
    }

    void ElfFile::addShdr(const std::string &sh_name, Elf32_Off sh_offset, Elf32_Word sh_size) {
        u32 offset = sizeof(Elf32_Ehdr) + sh_offset;      /* Section file offset */
        if (sh_name == ".text") {
            addShdr(sh_name, SHT_PROGBITS, SHF_ALLOC | SHF_EXECINSTR, 0, offset, sh_size, 0, 0, 4, 0);
        } else if (sh_name == ".data") {
            addShdr(sh_name, SHT_PROGBITS, SHF_ALLOC | SHF_WRITE, 0, offset, sh_size, 0, 0, 4, 0);
        } else if (sh_name == ".bss") {
            addShdr(sh_name, SHT_NOBITS, SHF_ALLOC | SHF_WRITE, 0, offset, sh_size, 0, 0, 4, 0);
        }
    }

    u32 ElfFile::getSegIndex(const std::string &segName) {
        u32 index = 0;
        for (const auto &shdrName: shdrNames) {
            if (shdrName == segName)//找到段
                break;
            ++index;
        }
        return index;
    }

    std::shared_ptr<RelocationItem> ElfFile::addRel(const std::string &segName, u32 addr,
                                                    const std::string &lb, u32 type) {
        std::shared_ptr<RelocationItem> relocationItem = std::make_shared<RelocationItem>();
        relocationItem->segName = segName;
        relocationItem->rel->r_offset = addr;       /* Address */
        relocationItem->rel->r_info = type;         /* Relocation type and symbol index */
        relocationItem->relName = lb;
        relTab.push_back(relocationItem);
        return relocationItem;
    }

    void ElfFile::assemblyObj() {
        std::vector<std::string> allSegNames = shdrNames;
        allSegNames.emplace_back(".shstrtab");  // 段表字符串表
        allSegNames.emplace_back(".symtab");    // 符号表
        allSegNames.emplace_back(".strtab");    // 字符串表
        allSegNames.emplace_back(".rel.text");  // 重定位代码段表
        allSegNames.emplace_back(".rel.data");  // 重定位数据段表

        // 段索引
        std::unordered_map<std::string, u32> shIndex;
        // 段名索引
        std::unordered_map<std::string, u32> shstrIndex;
        // 建立索引
        u32 segNum = allSegNames.size();
        for (u32 i = 0; i < segNum; ++i) {
            std::string name = allSegNames[i];
            shIndex[name] = i;
            shstrIndex[name] = shstrtab.size();
            shstrtab += name;
            shstrtab.push_back('\0');
        }

        // 符号索引
        std::unordered_map<std::string, u32> symIndex;
        // 符号名索引
        std::unordered_map<std::string, u32> strIndex;
        // 建立索引
        u32 symNum = symNames.size();
        for (u32 i = 0; i < symNum; ++i) {
            std::string name = symNames[i];
            symIndex[name] = i;
            strIndex[name] = strtab.size();
            strtab += name;
            strtab.push_back('\0');
        }

        // 更新符号表符号名索引
        for (u32 i = 0; i < symNum; ++i) {
            std::string name = symNames[i];
            symTab[name]->st_name = strIndex[name];
        }

        // 处理重定位表
        u32 relNum = relTab.size();
        for (u32 i = 0; i < relNum; ++i) {
            std::shared_ptr<Elf32_Rel> rel = std::make_shared<Elf32_Rel>();
            rel->r_offset = relTab[i]->rel->r_offset;   /* Address */
            rel->r_info = ELF32_R_INFO(symIndex[relTab[i]->relName],
                                       ELF32_R_TYPE(relTab[i]->rel->r_info));  /* Relocation type and symbol index */

            if (relTab[i]->segName == ".text") {
                relTextTab.emplace_back(rel);
            } else if (relTab[i]->segName == ".data") {
                relDataTab.emplace_back(rel);
            }
        }

        // 处理文件头
        char magic[] = {
                0x7f, 0x45, 0x4c, 0x46, 0x01, 0x01, 0x01, 0x00,
                0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
        };

        memcpy(&ehdr.e_ident, magic, sizeof(magic));
        ehdr.e_type = ET_REL;
        ehdr.e_machine = EM_386;
        ehdr.e_version = EV_CURRENT;
        ehdr.e_entry = 0;
        ehdr.e_phoff = 0;
        ehdr.e_shoff = 0;
        ehdr.e_flags = 0;
        ehdr.e_ehsize = sizeof(Elf32_Ehdr);
        ehdr.e_phentsize = 0;
        ehdr.e_phnum = 0;
        ehdr.e_shentsize = sizeof(Elf32_Shdr);
        ehdr.e_shnum = allSegNames.size();
        ehdr.e_shstrndx = shIndex[".shstrtab"];

        u32 curOff = sizeof(ehdr);
        for (auto &shdr: shdrTab) {
            curOff += shdr.second->sh_size;
            curOff += (4 - curOff % 4) % 4;
        }

        addShdr(".shstrtab", SHT_STRTAB, 0, 0, curOff, shstrtab.size(), SHN_UNDEF, 0, 1, 0);
        curOff += shstrtab.size();
        curOff += (4 - curOff % 4) % 4;

        ehdr.e_shoff = curOff;
        curOff += ehdr.e_shnum * ehdr.e_shentsize;

        addShdr(".symtab", SHT_SYMTAB, 0, 0, curOff, symNames.size() * sizeof(Elf32_Sym), shIndex[".strtab"], 0, 1,
                sizeof(Elf32_Sym));
        curOff += symNames.size() * sizeof(Elf32_Sym);

        addShdr(".strtab", SHT_STRTAB, 0, 0, curOff, strtab.size(), SHN_UNDEF, 0, 1, 0);
        curOff += strtab.size();
        curOff += (4 - curOff % 4) % 4;

        addShdr(".rel.text", SHT_REL, 0, 0, curOff, relTextTab.size() * sizeof(Elf32_Rel), shIndex[".symtab"],
                shIndex[".text"], 1,
                sizeof(Elf32_Rel));
        curOff += relTextTab.size() * sizeof(Elf32_Rel);

        addShdr(".rel.data", SHT_REL, 0, 0, curOff, relDataTab.size() * sizeof(Elf32_Rel), shIndex[".symtab"],
                shIndex[".data"], 1,
                sizeof(Elf32_Rel));
        curOff += relDataTab.size() * sizeof(Elf32_Rel);

        for (auto &name: allSegNames) {
            shdrTab[name]->sh_name = shstrIndex[name];
        }

    }

    void ElfFile::writeBytes(const std::string &sectionName, const std::vector<u8> &data) {
        for (auto &byte: data) {
            shdrData[sectionName].emplace_back(byte);
        }
    }


    void ElfFile::writeElf() {
        auto fileOut = fopen((m_fileName + ".o").c_str(), "wb");//输出文件
        assemblyObj();//组装文件
        fwrite(&ehdr, ehdr.e_ehsize, 1, fileOut);  //输出elf文件头

        //输出.text
        char buffer[1024] = {0};
        u32 numOfByte = -1;
        while (numOfByte != 0) {
            numOfByte = readShdrData(".text", buffer, 1024);
            fwrite(buffer, 1, numOfByte, fileOut);
        }

        //cout<<"---------------------输出指令的长度="<<inLen<<"  计算指令长度="<<shdrTab[".text"]->sh_size<<endl;
        padSeg(".text", ".data", fileOut);
        //.bss不用输出，对齐即可
        numOfByte = -1;
        while (numOfByte != 0) {
            numOfByte = readShdrData(".data", buffer, 1024);
            fwrite(buffer, 1, numOfByte, fileOut);
        }
        padSeg(".data", ".bss", fileOut);
        //.shstrtab，段表，符号表，.strtab，.rel.text，.rel.data
        writeElfTail(fileOut);//输出剩下的尾部
        fclose(fileOut);
    }


    void ElfFile::padSeg(const std::string &first, const std::string &second, FILE *fileOut) {
        char pad[1] = {0};
        u32 padNum = shdrTab[second]->sh_offset - (shdrTab[first]->sh_offset + shdrTab[first]->sh_size);
        while (padNum--) {
            fwrite(pad, 1, 1, fileOut);  //填充
        }
    }

    void ElfFile::writeElfTail(FILE *fileOut) {
        //-----输出
        fwrite(shstrtab.c_str(), shstrtab.size(), 1, fileOut); //.shstrtab

        char pad[1] = {0};
        u32 padNum = ehdr.e_shoff - (shdrTab[".shstrtab"]->sh_offset + shdrTab[".shstrtab"]->sh_size);
        while (padNum--) {
            fwrite(pad, 1, 1, fileOut);  //填充
        }

        for (const auto &shdrName: shdrNames) {
            //段表
            auto shdr = shdrTab[shdrName];
            fwrite(reinterpret_cast<char *>(shdr.get()), ehdr.e_shentsize, 1, fileOut);
        }

        for (const auto &symName: symNames) {
            //符号表
            auto sym = symTab[symName];
            fwrite(sym.get(), sizeof(Elf32_Sym), 1, fileOut);
        }

        //.strtab
        fwrite(strtab.c_str(), strtab.size(), 1, fileOut);
        padSeg(".strtab", ".rel.text", fileOut);
        for (auto &rel: relTextTab) {
            //.rel.text
            fwrite(rel.get(), sizeof(Elf32_Rel), 1, fileOut);
        }
        for (auto &rel: relDataTab) {
            //.rel.data
            fwrite(rel.get(), sizeof(Elf32_Rel), 1, fileOut);
        }
    }

    u32 ElfFile::readShdrData(const std::string &segName, char *buffer, u32 bufferSize) {
        u32 numOfByte = 0;
        auto bytes = shdrData[segName];
        for (auto &byte: bytes) {
            buffer[numOfByte] = static_cast<char>(byte);
            ++numOfByte;
            if (numOfByte == bufferSize) {
                break;
            }
        }
        for (u32 i = 0; i < numOfByte; ++i) {
            shdrData[segName].pop_front();
        }
        return numOfByte;
    }
}